import ddddocr
from PIL import Image
from os import remove


def get_yzm(n=4):
    ocr = ddddocr.DdddOcr()
    with open('captcha.jpg', 'rb') as f:
        img_bytes = f.read()
    res = ocr.classification(img_bytes)
    if len(res) != n:
        return False
    else:
        remove('captcha.jpg')
        return res


def handle_yzm(driver, captcha_element):
    # 获取验证码图片的坐标和大小
    location = captcha_element.location
    size = captcha_element.size
    captcha_path = 'captcha.jpg'
    driver.save_screenshot(captcha_path)
    captcha_image = Image.open(captcha_path)
    captcha_image = captcha_image.crop(
        (location['x'], location['y'], location['x'] + size['width'], location['y'] + size['height']))
    captcha_image = captcha_image.convert('RGB')
    captcha_image.save(captcha_path)
    return get_yzm()

